"""
将分析结果上传到PG数据库中
"""
import psycopg2
import psycopg2.extras
from typing import Literal, Union, Any
from core.utils.path import result_dir
import pickle
from pathlib import Path
import json


def pretty(obj: Any):
    """美化显示"""
    if isinstance(obj, int):
        return obj
    elif isinstance(obj, dict) or isinstance(obj, list):
        return json.dumps(obj, indent=4, ensure_ascii=False)
    else:
        return str(obj)


class TableEngine(object):

    def __init__(self):
        self.conn = psycopg2.connect(
            database="dianda_spreadsheet",
            user="root",
            password="00123.Com",
            host="172.18.0.66",
            port="5400"
        )
        self.table_name = "novelai"

    def execute(self, sql: str, return_type: Literal["none", "one", "all"] = "none"):
        """sql 执行"""
        cursor = self.conn.cursor()
        cursor.execute(sql)
        if return_type == "one":
            data = cursor.fetchone()
        elif return_type == "all":
            data = cursor.fetchall()
        else:
            data = None
        self.conn.commit()
        return data

    def create_novel_table(self):
        """创建表"""
        sql = '''CREATE TABLE NovelAnalyze(
            id         SERIAL PRIMARY KEY     NOT NULL,
            model             TEXT            NOT NULL,
            title             TEXT            NOT NULL,
            topn              INT             NOT NULL,
            characters        TEXT            NOT NULL,
            plot_rating       TEXT            NOT NULL,
            character_rating  TEXT            NOT NULL,
            worldview_rating  TEXT            NOT NULL,
            style_rating      TEXT            NOT NULL,
            total_score       TEXT            NOT NULL
        );'''
        self.execute(sql)

    def create_chapter_table(self):
        """创建表"""
        sql = '''CREATE TABLE NovelChapters(
            id         SERIAL PRIMARY KEY     NOT NULL,
            model             TEXT            NOT NULL,
            novel_id          INT             NOT NULL,
            novel_title       TEXT            NOT NULL,
            volume            TEXT            NOT NULL,
            chapter_number    INT             NOT NULL,
            chapter_title     TEXT            NOT NULL,
            chapter_content   TEXT            NOT NULL,
            character_profile TEXT            NOT NULL,
            characters        TEXT            NOT NULL, 
            plot              TEXT            NOT NULL,
            characterization  TEXT            NOT NULL,
            worldview         TEXT            NOT NULL, 
            writing_style     TEXT            NOT NULL,
            problems          TEXT            NOT NULL
        );'''
        self.execute(sql)

    def show_tables(self):
        sql = "select table_name from information_schema.tables where table_schema='public';"
        data = engine.execute(sql, return_type="all")
        print(data)

    def insert_novel(self, novel_file: Union[Path, str], model: str = "doubao-seed-1-6-250615-thinking"):
        """插入小说数据和小说的章节数据"""
        with open(novel_file, 'rb') as fs:
            novel = pickle.load(fs)

        cursor = self.conn.cursor()

        # 插入小说数据
        query = ("INSERT INTO novelanalyze (model,title,topn,characters,plot_rating,"
                 "character_rating,worldview_rating,style_rating,total_score) VALUES %s")
        novel_data = [pretty(e) for e in [
            model, novel.title, novel.topn, novel.characters, novel.plot_rating,
            novel.character_rating, novel.worldview_rating, novel.style_rating, novel.total_score
        ]]
        psycopg2.extras.execute_values(cursor, query, [novel_data])

        # 查询 novel 对应的 id
        sql = f"select * from novelanalyze where title='{novel.title}'"
        datas = self.execute(sql, return_type="all")
        if len(datas) == 0:
            return
        novel_id = datas[-1][0]

        query = (
            "INSERT INTO novelchapters (model,novel_id,novel_title,volume,chapter_number,chapter_title,chapter_content,"
            "character_profile,characters,plot,characterization,worldview, writing_style, problems) VALUES %s")
        chapter_datas = []
        for chapter in novel.chapters:
            chapter_data = [pretty(e) for e in [
                model,
                novel_id, novel.title, chapter.volume, chapter.chapter_number, chapter.chapter_title,
                chapter.chapter_content, chapter.character_profile, chapter.characters,
                chapter.plot, chapter.characterization, chapter.worldview, chapter.writing_style, chapter.problems
            ]]
            chapter_datas.append(chapter_data)
        psycopg2.extras.execute_values(cursor, query, chapter_datas)

    def show_novels(self):
        datas = self.execute("select * from novelanalyze", return_type="all")
        for data in datas:
            print(data)

    def show_chapters(self):
        datas = self.execute("select * from novelchapters", return_type="all")
        for data in datas:
            print(data)

    def delete_tables(self):
        sql = f"DROP TABLE IF EXISTS novelanalyze;"
        self.execute(sql, return_type="none")
        sql = f"DROP TABLE IF EXISTS novelchapters;"
        self.execute(sql, return_type="none")

    def delete_records(self):
        sql = f"delete from novelanalyze where id > 0"
        self.execute(sql, return_type="none")
        sql = f"delete from novelchapters where id > 0"
        self.execute(sql, return_type="none")


if __name__ == '__main__':

    engine = TableEngine()
    # engine.delete_records()
    engine.delete_tables()
    engine.create_novel_table()
    engine.create_chapter_table()
    engine.show_tables()

    # 上传分析结果
    pkl_dir = result_dir
    for novel_file in pkl_dir.rglob('*(100章).pkl'):
        print(novel_file)
        engine.insert_novel(novel_file, model="doubao-seed-1-6-250615-thinking")

    engine.show_novels()
    engine.show_chapters()
